package com.zelix.analy;

import com.arthas.client.ArthasClient;
import com.arthas.client.ArthasTool;
import com.zelix.tool.MethodConvert;
import org.objectweb.asm.ClassReader;
import org.objectweb.asm.Handle;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.tree.*;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Enumeration;
import java.util.List;
import java.util.jar.JarEntry;
import java.util.jar.JarFile;
import java.util.zip.ZipEntry;

public class analyJetbrainsAgent {
    public static final String STRING = "(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;";
    public static final String METHOD = "(Ljava/lang/Object;Ljava/lang/Object;Ljava/lang/Object;)Ljava/lang/Object;";

    public static void main(String[] args) throws IOException {
        JarFile jarFile = new JarFile("C:\\Users\\googlewell\\Desktop\\jetbrains-agent.jar");
        String rule = "fuck_the_regulations_v320/bL";

        ArthasClient client = ArthasTool.getClient();

        Enumeration<JarEntry> entries = jarFile.entries();
        while (entries.hasMoreElements()) {
            JarEntry entry = entries.nextElement();
            if (entry.getName().endsWith(".class") && entry.getName().startsWith(rule)) {//只分析class文件
                test(client, jarFile, entry.getName().substring(0, entry.getName().length() - 6));
            }
        }

        client.disconnect();
    }

    public static void test(ArthasClient client, JarFile jarFile, String cls) throws IOException {
        String classLoader = ArthasTool.getClassLoader(client, cls);
        if (classLoader == null) {
            System.out.println("类 " + cls + " 还未加载,请点击相应功能触发其加载.");
            return;
        }

        ZipEntry entry = jarFile.getEntry(cls + ".class");
        //字符串加密的ognl表示方式
        String stringMethod = null;

        if (entry != null) {
            ClassReader cr = new ClassReader(jarFile.getInputStream(entry));
            ClassNode cn = new ClassNode();
            cr.accept(cn, 0);

            List<MethodNode> methods = cn.methods;

            //查找符合的方法并初始化解密方法
            for (MethodNode methodNode : methods) {
                if (methodNode.desc.equals("(IJ)Ljava/lang/String;")) {
                    stringMethod = "@" + cls.replaceAll("/", ".") + "@" + methodNode.name;
                }
            }

            for (MethodNode methodNode : methods) {

                System.out.println("-------------(" + methodNode.name + ":" + methodNode.desc + ")-------------");
                InsnList insnList = methodNode.instructions;

                //计算long值,遍历每个方法前初始化,不同方法值不同
                long l = 0;

                //遍历指令
                for (int i = 0; i < insnList.size(); i++) {
                    /**
                     * GETSTATIC fuck_the_regulations_v320/bL.b J
                     * LDC 17435732912823L
                     * LXOR
                     * LSTORE 0
                     */
                    if (insnList.get(i) instanceof VarInsnNode && i >= 3) {
                        VarInsnNode varInsnNode = (VarInsnNode) insnList.get(i);
                        if (varInsnNode.getOpcode() == Opcodes.LSTORE && insnList.get(i - 1) instanceof InsnNode && insnList.get(i - 2) instanceof LdcInsnNode && insnList.get(i - 3) instanceof FieldInsnNode) {
                            InsnNode insnNode = (InsnNode) insnList.get(i - 1);
                            LdcInsnNode ldcInsnNode = (LdcInsnNode) insnList.get(i - 2);
                            FieldInsnNode fieldInsnNode = (FieldInsnNode) insnList.get(i - 3);
                            if (insnNode.getOpcode() == Opcodes.LXOR && fieldInsnNode.getOpcode() == Opcodes.GETSTATIC && fieldInsnNode.desc.equals("J")) {
                                String rtn = client.sendCommand("ognl '@" + fieldInsnNode.owner.replaceAll("/", ".") + "@" + fieldInsnNode.name + ".toString()' -c " + classLoader);
                                rtn = ArthasTool.removeStringLebal(rtn);
                                l = Long.valueOf(ldcInsnNode.cst.toString()) ^ Long.valueOf(rtn);
                            }
                        }
                    }
                    if (insnList.get(i) instanceof InvokeDynamicInsnNode) {
                        InvokeDynamicInsnNode dynamicNode = (InvokeDynamicInsnNode) insnList.get(i);
                        System.out.println("------------------");
                        Handle bsm = dynamicNode.bsm;
                        String out = bsm.getOwner() + ":" + bsm.getName() + ":" + bsm.getDesc();
                        if (bsm.getDesc().equals(STRING)) {
                            if (dynamicNode.desc.equals("(IJ)Ljava/lang/String;")) {
                                /**
                                 * SIPUSH 27134
                                 * LDC 7612129396686384391L
                                 * LLOAD 1
                                 * LXOR
                                 * INVOKEDYNAMIC r (IJ)Ljava/lang/String; handle[H_INVOKESTATIC com/zelix/ZKM.b(Ljava/lang/invoke/MethodHandles$Lookup;Ljava/lang/String;Ljava/lang/invoke/MethodType;)Ljava/lang/invoke/CallSite;] args[]
                                 * 完全匹配上述模式,如果不符,暂不处理.可能存在该加密方法,但是被标签分割,不修改,修改可能导致代码错乱
                                 */
                                if (insnList.get(i - 1) instanceof InsnNode && insnList.get(i - 1).getOpcode() == Opcodes.LXOR) {//判断上一个是不是LXOR
                                    if (insnList.get(i - 2) instanceof VarInsnNode && insnList.get(i - 2).getOpcode() == Opcodes.LLOAD) {
                                        if (insnList.get(i - 3) instanceof LdcInsnNode) {
                                            if (insnList.get(i - 4) instanceof IntInsnNode) {
                                                LdcInsnNode ldcNode = (LdcInsnNode) insnList.get(i - 3);
                                                IntInsnNode intNode = (IntInsnNode) insnList.get(i - 4);

                                                if (l != 0) {
                                                    String command = "ognl '" + stringMethod + "(" + intNode.operand + "," + ldcNode.cst + "L^" + l + "L)' -c " + classLoader;
                                                    String rtn = client.sendCommand(command);
                                                    rtn = ArthasTool.removeStringLebal(rtn);
                                                    System.out.println("字符串解密:" + rtn);
                                                } else {
                                                    System.out.println("字符串解密:key值未初始化");
                                                }
                                            }
                                        }
                                    }
                                }
                            } else {
                                System.out.println("未知字符串解密:" + out);
                            }
                        } else if (bsm.getDesc().equals(METHOD)) {
                            //这里的解析类固定,可以解析out的bsm方法,遍历MethodInsnNode,获取方法返回参数是Ljava/lang/reflect/Method;的方法,这基本就是实现方法
                            String rtn = client.sendCommand("ognl '@fuck_the_regulations_v320.us@q(" + Integer.valueOf(dynamicNode.name, 32) + ").toString()' -c " + classLoader);
                            rtn = ArthasTool.removeStringLebal(rtn);
                            System.out.println("调用方法:" + rtn);
                            MethodConvert.analyMethod(rtn);
                        } else {
                            System.out.println("未知操作:" + out);
                            System.exit(0);
                        }
                    }

                    /**
                     * LDC "驃岂Ⲅ✾籮⫑ꦞ暐》縖莜㈜褦작폵␚뒢憥흌汱㜾"
                     * INVOKEVIRTUAL java/lang/String.toCharArray()[C
                     * DUP
                     * DUP
                     * BIPUSH 9
                     * DUP_X1
                     * CALOAD
                     * SIPUSH 253
                     * IXOR
                     * I2C
                     * CASTORE
                     * SIPUSH 28296
                     * ICONST_0
                     * ICONST_0
                     * INVOKESTATIC fuck_the_regulations_v320/di$cm.m(Ljava/lang/Object;SZZ)Ljava/lang/String;
                     */
                    if (insnList.get(i) instanceof MethodInsnNode && i >= 14) {
                        MethodInsnNode methodInsnNode = (MethodInsnNode) insnList.get(i);
                        if (methodInsnNode.getOpcode() == Opcodes.INVOKESTATIC && methodInsnNode.desc.equals("(Ljava/lang/Object;SZZ)Ljava/lang/String;")) {

                            /**
                             //前两个bool
                             if (insnList.get(i - 1) instanceof InsnNode && insnList.get(i - 2) instanceof InsnNode) {
                             InsnNode insnNode1 = (InsnNode) insnList.get(i - 1);
                             InsnNode insnNode2 = (InsnNode) insnList.get(i - 2);
                             if ((insnNode1.getOpcode() == Opcodes.ICONST_0 || insnNode1.getOpcode() == Opcodes.ICONST_1) && (insnNode2.getOpcode() == Opcodes.ICONST_0 || insnNode2.getOpcode() == Opcodes.ICONST_1)) {//bool值，只能是0和1
                             int arg3 = getInsnNodeValue(insnNode2.getOpcode());
                             int arg4 = getInsnNodeValue(insnNode1.getOpcode());
                             //第三个short
                             if (insnList.get(i - 3) instanceof IntInsnNode) {
                             IntInsnNode intInsnNode = (IntInsnNode) insnList.get(i - 3);
                             if (intInsnNode.getOpcode() == Opcodes.SIPUSH) {
                             int arg2 = intInsnNode.operand;
                             System.out.println(arg2);
                             }
                             }
                             }
                             }
                             **/

                            //不再详细判断数值和Opcodes,只判断类型,所有类型都匹配概率很小
                            if (insnList.get(i - 1) instanceof InsnNode && insnList.get(i - 2) instanceof InsnNode//两个bool类型值
                                    && insnList.get(i - 3) instanceof IntInsnNode//SIPUSH
                                    && insnList.get(i - 4) instanceof InsnNode && insnList.get(i - 5) instanceof InsnNode && insnList.get(i - 6) instanceof InsnNode//CASTORE I2C IXOR
                                    && insnList.get(i - 7) instanceof IntInsnNode//SIPUSH
                                    && insnList.get(i - 8) instanceof InsnNode && insnList.get(i - 9) instanceof InsnNode//CALOAD DUP_X1
                                    && (insnList.get(i - 10) instanceof InsnNode || insnList.get(i - 10) instanceof IntInsnNode)//ICONST_2 or BIPUSH
                                    && insnList.get(i - 11) instanceof InsnNode && insnList.get(i - 12) instanceof InsnNode//DUP DUP
                                    && insnList.get(i - 13) instanceof MethodInsnNode//INVOKEVIRTUAL java/lang/String.toCharArray()[C
                                    && insnList.get(i - 14) instanceof LdcInsnNode//LDC
                            ) {
                                InsnNode insnNode1 = (InsnNode) insnList.get(i - 1);
                                InsnNode insnNode2 = (InsnNode) insnList.get(i - 2);
                                int arg3 = getInsnNodeValue(insnNode2.getOpcode());
                                int arg4 = getInsnNodeValue(insnNode1.getOpcode());

                                IntInsnNode intInsnNode = (IntInsnNode) insnList.get(i - 3);
                                int arg2 = intInsnNode.operand;

                                intInsnNode = (IntInsnNode) insnList.get(i - 7);
                                int xorint = intInsnNode.operand;

                                int index = -1;
                                if (insnList.get(i - 10) instanceof InsnNode) {
                                    insnNode2 = (InsnNode) insnList.get(i - 10);
                                    index = getInsnNodeValue(insnNode2.getOpcode());
                                } else {
                                    intInsnNode = (IntInsnNode) insnList.get(i - 10);
                                    index = intInsnNode.operand;
                                }

                                LdcInsnNode ldcInsnNode = (LdcInsnNode) insnList.get(i - 14);
                                char[] chars = ldcInsnNode.cst.toString().toCharArray();
                                chars[index] = (char) (chars[index] ^ xorint);
                                String base64 = Base64.getEncoder().encodeToString(String.valueOf(chars).getBytes(StandardCharsets.UTF_8));
                                String cmd = "ognl '#bytes=@java.util.Base64@getDecoder().decode(\"" + base64 + "\"),#str=new String(#bytes,\"utf-8\"),#rtn=@" + methodInsnNode.owner.replaceAll("/", ".") + "@" + methodInsnNode.name + "(#str.toCharArray()," + arg2 + "," + arg3 + "," + arg4 + "),@java.util.Base64@getEncoder().encodeToString(#rtn.getBytes(\"utf-8\"))'";
                                String rtn = client.sendCommand(cmd);
                                rtn = ArthasTool.removeStringLebal(rtn);
                                System.out.println(cmd);
                                System.out.println("字符串解密:" + new String(Base64.getDecoder().decode(rtn), StandardCharsets.UTF_8));
                            }
                        }
                    }
                }


                //写入文件
                //ClassWriter cw = new ClassWriter(ClassWriter.COMPUTE_MAXS);
                //cn.accept(cw);
                //AnalyTool.classToFile("f.class", cw.toByteArray());
            }
        }
    }

    private static int getInsnNodeValue(int type) {
        switch (type) {
            case Opcodes.ICONST_0:
                return 0;
            case Opcodes.ICONST_1:
                return 1;
            case Opcodes.ICONST_2:
                return 2;
            case Opcodes.ICONST_3:
                return 3;
            case Opcodes.ICONST_4:
                return 4;
            case Opcodes.ICONST_5:
                return 5;
            case Opcodes.ICONST_M1:
                return -1;
            default:
                return -2;
        }
    }
}
